import numpy as np
import openpyxl


def calculate_weight_matrix(size):
    w, h = size
    x = np.linspace(0, 1, w)
    y = np.linspace(0, 1, h)
    xv, yv = np.meshgrid(x, y)
    weight_matrix = xv * yv
    return weight_matrix


def save_weight_matrix_to_excel(weight_matrix, sheet, start_row=1, start_col=1):
    for i in range(weight_matrix.shape[0]):
        for j in range(weight_matrix.shape[1]):
            sheet.cell(row=start_row + i, column=start_col + j, value=weight_matrix[i, j])


def save_all_weight_matrices_to_excel(output_excel, size):
    workbook = openpyxl.Workbook()

    # 计算一个角的权重矩阵
    weight_matrix = calculate_weight_matrix((size, size))

    # 计算并保存四个角的权重矩阵
    sheets_info = {
        'Top_Left': weight_matrix,
        'Top_Right': np.fliplr(weight_matrix),
        'Bottom_Left': np.flipud(weight_matrix),
        'Bottom_Right': np.flipud(np.fliplr(weight_matrix)),
    }

    # 计算并保存总权重矩阵
    total_weight_matrix = sum(sheets_info.values())
    sheet = workbook.create_sheet(title='Total_Weight')
    save_weight_matrix_to_excel(total_weight_matrix, sheet)

    # 计算并保存额外的权重矩阵和
    combined_weights_info = {
        'Top_Left_Right_Top': sheets_info['Top_Left'] + sheets_info['Top_Right'],
        'Top_Left_Bottom_Left': sheets_info['Top_Left'] + sheets_info['Bottom_Left'],
        'Right_Top_Right_Bottom': sheets_info['Top_Right'] + sheets_info['Bottom_Right'],
        'Bottom_Left_Bottom_Right': sheets_info['Bottom_Left'] + sheets_info['Bottom_Right']
    }

    for sheet_name, weight_matrix in combined_weights_info.items():
        sheet = workbook.create_sheet(title=sheet_name)
        save_weight_matrix_to_excel(weight_matrix, sheet)

    # 保存单独的权重矩阵
    for sheet_name, weight_matrix in sheets_info.items():
        sheet = workbook.create_sheet(title=sheet_name)
        save_weight_matrix_to_excel(weight_matrix, sheet)

    workbook.remove(workbook['Sheet'])  # 移除默认创建的Sheet
    workbook.save(output_excel)
    print(f"Saved all weight matrices to {output_excel}")
# 示例用法
size = 256
output_excel = 'E:\\wafer52\\weight\\path_to_all_weight_matrices.xlsx'
save_all_weight_matrices_to_excel(output_excel, size)
